from paddleocr import PPDocTranslation
import time
import sys

print("开始导入模块...")
# Create a translation pipeline
print("初始化PPDocTranslation...")
pipeline = PPDocTranslation()
print("PPDocTranslation初始化完成！")

# Document path
input_path = "input_files/DeepSeek-V3-2pages.pdf"
# input_path = "input_files/en-1.md"
# input_path="input_files/DeepSeek-V3-2pages_1747906036.md"

# Output directory
output_path = "./output"

print(f"开始处理文件：{input_path}")

if input_path.lower().endswith(".md"):
    # Read markdown documents, supporting passing in directories and url links with the .md suffix
    print("开始加载Markdown文件...")
    ori_md_info_list = pipeline.load_from_markdown(input_path)
    print("Markdown文件加载完成！")
else:
    # Use PP-StructureV3 to perform layout parsing on PDF/image documents to obtain markdown information
    print("开始使用PP-StructureV3进行布局分析...")
    try:
        visual_predict_res = pipeline.visual_predict(
            input_path,
            use_doc_orientation_classify=False,
            use_doc_unwarping=False,
            use_common_ocr=True,
            use_seal_recognition=True,
            use_table_recognition=True,
        )
        print("PP-StructureV3布局分析完成！")
        
        ori_md_info_list = []
        for i, res in enumerate(visual_predict_res):
            print(f"处理第 {i+1}/{len(visual_predict_res)} 页...")
            layout_parsing_result = res["layout_parsing_result"]
            ori_md_info_list.append(layout_parsing_result.markdown)
            print(f"保存第 {i+1} 页图片...")
            layout_parsing_result.save_to_img(output_path)
            print(f"保存第 {i+1} 页Markdown...")
            layout_parsing_result.save_to_markdown(output_path)
            print(f"第 {i+1} 页处理完成")
        
        # Concatenate the markdown information of multi-page documents into a single markdown file, and save the merged original markdown text
        if input_path.lower().endswith(".pdf"):
            print("合并多页文档的Markdown信息...")
            ori_md_info = pipeline.concatenate_markdown_pages(ori_md_info_list)
            print("保存合并后的原始Markdown文本...")
            ori_md_info.save_to_markdown(output_path)
            print("保存完成！")
            
        print("所有页面处理完毕，无需翻译，程序结束。")
            
    except Exception as e:
        print(f"处理过程中出错：{str(e)}")
        import traceback
        traceback.print_exc()
